
class Unit:
    def __init__(self, unit, utype):
        self.uid = int(unit['ID'])
        self.side = 'red' if unit['JB'] else 'blue'
        self.type = utype
        self.name = unit['XH']
        self.speed = unit.get('SP', 10)
        self.isalive = unit.get('WH', 1)#['WH']
        self.health = unit.get('DA', 100)
        try:
            self.x = round(unit['X'] - 483e4, 3)
            self.y = round(unit['Y'] - 755e4, 3)
            self.z = round(unit['Z'], 0)
        #self.state = unit['ST']
        #self.weapons = unit['WP']
        except Exception as e:
            print('----------Error init-------------')
            print(f'Exception : {e}')
            print(unit, utype)
            self.x, self.y, self.z = -1, -1, -1
            self.isalive = 0
            print('---------------------------------')

class Side:
    def __init__(self, info, obs):
        self.__dict__ = {"tanks": [], 'wrjs': [], 'bbb':[], 'wrc': [], 'yzc': [], 'bgc': [], 'sb': [], 'obs': obs}
        self.__dict__.update({"qb_tanks": [], 'qb_wrjs': [], 'qb_bbb':[], 'qb_wrc': [], 'qb_yzc': [], 'qb_bgc': [], 'qbs':[], 'qb_sb': []})
        self.load(info)

    def load(self, info):
        for unit in info['units']:
            if unit['XH'] in ['T96A', 'M2A3', 'M1A2']:
                u = Unit(unit, '坦克')
                self.tanks.append(u)
            elif '无人机' in unit['XH'] :
                u = Unit(unit, '无人机')
                self.wrjs.append(u)
            elif '04A' == unit['XH']:
                u = Unit(unit, '步兵班')
                self.bbb.append(u)
            elif '无人车' in unit['XH']:
                u = Unit(unit, '无人车')
                self.wrc.append(u)
            elif '运载车' in unit['XH']:
                u = Unit(unit, '运载车')
                self.yzc.append(u)
            elif '补给车' in unit['XH']:
                u = Unit(unit, '补给车')
                self.bgc.append(u)
            #elif '战壕' in unit['XH']:
                #u = Unit(unit, '战壕')
                #self.sb.append(u)
            elif '士兵' in unit['XH']:
                u = Unit(unit, '士兵')
                self.sb.append(u)
            self.obs.unit_dict[u.uid] = u
        for unit in info['qb']:
            if unit['XH'] in ['T96A', 'M2A3', 'M1A2']:
                u = Unit(unit, '坦克')
                self.qb_tanks.append(u)
            elif '无人机' in unit['XH'] :
                u = Unit(unit, '无人机')
                self.qb_wrjs.append(u)
            elif '' == unit['XH']:
                u = Unit(unit, '步兵班')
                self.qb_bbb.append(u)
            elif '无人车' in unit['XH']:
                u = Unit(unit, '无人车')
                self.qb_wrc.append(u)
            #elif '运载车' in unit['XH']:
            #    u = Unit(unit, '运载车')
            #    self.qb_yzc.append(u)
            #elif '补给车' in unit['XH']:
            #    u = Unit(unit, '补给车')
            #    self.qb_bgc.append(u)
            elif '士兵' in unit['XH']:
                u = Unit(unit, '士兵')
                self.qb_sb.append(u)
            self.qbs.append(u)


class Observation:
    def __init__(self, state):
        self.state = state
        self.sim_time = state['sim_time']
        self.unit_dict = {}
        self.red = Side(state['red'], self)
        self.blue = Side(state['blue'], self)

    def print_state(self):
        print("SimTime: ", self.state['sim_time'])
        print('----'*20)
        for team in ['red', 'blue']:
            #continue
            print(f'{team} Observation: \nunits:')
            runits = self.state[team]['units']
            runits.sort(key=lambda x:int(x['ID']))
            for unit in runits:
                #if int(unit['ID']) not in [85, 86, 87, 88, 89 ,90]:
                    #print(unit['ID'])
                #    continue
                #print(unit)
                if "战壕" in unit['XH']:  
                    continue
                print(f"\t{ 'Red 'if unit['JB'] else 'Blue'}\t{'坦克'+unit['XH'] if '96A' in unit['XH'] else unit['XH']}\t{unit['ID']}\t\t[{round(unit['X']-483e4, 3)}, {round(unit['Y']-755e4, 3)}, {int(unit['Z'])}] \t ---{unit.get('WH', 1)} --{unit.get('DA', 100)}")
            print('----'*20)
        return 
        units = list(self.unit_dict.values())
        units.sort(key=lambda x:int(x.uid))
        for side in ['Red', 'Blue']:
            print(f'{side} Observation: \nunits:')
            for unit in units:
                if '战壕' in unit.name:
                    continue
                if unit.side == side:
                    print(unit.__dict__)
            print('----'*20)

    def print_state_idx(self, idx):
        unit = self.unit_dict[idx]
        #print(f"\t{ 'Red 'if unit['JB'] else 'Blue'}\t{'坦克'+unit['XH'] if '96A' in unit['XH'] else unit['XH']}\t{unit['ID']}\t\t[{round(unit['X']-483e4, 2)}, {round(unit['Y']-755e4, 2)}, {int(unit['Z'])}] \t ---{unit.get('WH', 1)} --{unit.get('DA', 100)}")
        print(f"\t{unit.side}\t{unit.name}\t{unit.uid}\t\t {unit.x}, {unit.y}, {unit.z}, \t ---{unit.isalive} --{unit.health} --{unit.speed}")

if __name__ == '__main__':
    import os
    from client import EnvClient
    def run(host, port):
        env_client = EnvClient(0, 6100, 'registry.inspir.ai:5000/lzy-combat:1.1', '/home/lzy_scen.ntedt')
        env_client.reset()
        state = env_client.get_observation()
        obs = state#Observation(state)
        obs.print_state()

    #host = os.popen("cat /etc/resolv.conf |grep nameserver |awk '{print $2}'").read().strip()
    host = '127.0.0.1'
    port = 6101
    run(host, port)
    #run('localhost', 6101)
